from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier   # 估计也不需要引入
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score   # 用于检测模型准确度的，大概率不需要引入
import numpy as np
import pandas as pd
from sklearn.model_selection import cross_val_score
import matplotlib.pyplot as plt
from sklearn.tree import plot_tree

import requests,json,time,random,csv

from app.models import db, Product, User   # 引入数据库
import app.utils.market as market
#from app import create_app
#from app.config import BaseConfig


# 数据集包含  X -  股票代码、 最新价、 涨跌幅、 涨跌额、  成交量、
#                 code    price  change1 change2  deal
#                 振幅、 最高、 最低、 今开、 昨收、 市盈率、市净率
#                range  high  low  open  close   PE    PB
#
#           y -  是否持有
#               sustained


# stage 1. 生成训练数据集
# stage 2. 训练模型


class DecisionTreeTrainer:
    def __init__(self, max_depth=5, min_samples_leaf=5, min_samples_split=10):
        """
        初始化决策树训练器
        参数:
        max_depth: 决策树最大深度
        min_samples_leaf: 叶节点最小样本数
        min_samples_split: 节点分裂最小样本数
        """
        self.tree = DecisionTreeClassifier(
            max_depth=max_depth,
            min_samples_leaf=min_samples_leaf,
            min_samples_split=min_samples_split
        )

        # 创建一个标签编码器（用于将股票代码转换为数字）, 其他数据都直接转换为数字
        self.label_encoder_code = LabelEncoder()

        # 存储训练数据, 默认为空
        self.X = None
        self.y = None

        # 更新特征名称（与market返回字段对应）
        self.feature_names = ['code', 'price', 'change1', 'change2', 'deal', 'range', 'high', 'low', 'open', 'close', 'PE', 'PB']
        self.target_name = 'sustained'


    def _load_data(self, user_id):
        """
        生成训练数据集
        """
        # step 0. 准备存储特征和标签————————————————————————————————————
        features = []   # 里面要包含股票代码 用转换器转换
        target = []

        # step 1. 从用户的持仓获取训练数据集——————————————————————————
        user = User.query.get(user_id)
        if not user:
            raise ValueError(f"用户ID {user_id} 不存在")

        # user_codes = set(p.code for p in products)   # 存储用户持仓的数据表 但是未过滤基金

        # 获取用户持仓股票代码（过滤掉基金）
        user_stock_codes = set()
        for product in user.products:
            # 假设Product有type字段标识类型：'stock'表示股票，'fund'表示基金
            if product.type == '股票':
                user_stock_codes.add(product.code)

        # logger.info(f"用户持仓股票数量: {len(user_stock_codes)}")

        # user_products = set(market.get_latest_stock_info(c) for c in user_stock_codes)

        # 获取该用户持仓里面的所有数据集
        positive_count = 0
        for code in user_stock_codes:
            stock_info = market.get_latest_stock_info(code)
            if stock_info is None:
                print(f"无法获取持仓股票数据: {code}")
                continue  # 如果获取失败就跳过
            features.append([
                stock_info['股票代码'],
                stock_info['最新价'],    # 注意这里是1297，网页上是12.97
                stock_info['涨跌幅'],    # 注意这里是2001，网页上是20.01%
                stock_info['涨跌额'],    # 注意这里是282，网页上是2.82
                stock_info['成交量'],    # 注意这里是637574，网页上是‘63.76万’
                stock_info['振幅'],     # 注意这里是1732，网页上是17.32%
                stock_info['最高'],     # 注意这里是1691，网页上是16.91
                stock_info['最低'],     # 注意这里是1447，网页上是14.47
                stock_info['今开'],     # 注意这里是3000，网页上是30.00
                stock_info['昨收'],     # 注意这里是1208，网页上是12.08
                stock_info['市盈率'],    # 注意这里是-7366，网页上是-73.66
                stock_info['换手率']     # 注意这里是1784，网页上是17.84%
            ])
            target.append(1)  # 持有标记为 1, 因为数据库里面的肯定持有了
            positive_count += 1
        #print(f"已添加 {positive_count} 个正样本")

        # step 2. 从市场上爬取一定的训练数据集—————————————————————————— 只爬股票市场的数据 爬取几个随机页数

        negative_count = 0
        num_pages = random.randint(3, 5)  # 随机选择爬取的总页数
        pages = random.sample(range(1, 287), num_pages)  # 获取随机页码

        for page in pages:
            #print(f"爬取第 {page} 页市场数据...")

            timestamp = int(time.time() * 1000)
            callback = f"jQuery37101004698719727698_{timestamp}"
            url = f"https://push2.eastmoney.com/api/qt/clist/get?np=1&fltt=1&invt=2&cb={callback}&fs=m%3A0%2Bt%3A6%2Cm%3A0%2Bt%3A80%2Cm%3A1%2Bt%3A2%2Cm%3A1%2Bt%3A23%2Cm%3A0%2Bt%3A81%2Bs%3A2048&fields=f12%2Cf13%2Cf14%2Cf1%2Cf2%2Cf4%2Cf3%2Cf152%2Cf5%2Cf6%2Cf7%2Cf15%2Cf18%2Cf16%2Cf17%2Cf10%2Cf8%2Cf9%2Cf23&fid=f3&pn={page}&pz=20&po=1&dect=1&ut=fa5fd1943c7b386f172d6893dbfba10b&wbp2u=%7C0%7C0%7C0%7Cweb&_{timestamp}"
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
                "Referer": "https://www.eastmoney.com/",
            }
            response = requests.get(url, headers=headers)
            response.raise_for_status()  # 检查 HTTP 错误

            json_str = response.text.split('(', 1)[1].rsplit(')', 1)[0]
            res_dict = json.loads(json_str)

            # 关键检查：验证数据结构
            if not res_dict.get('data') or not isinstance(res_dict['data'].get('diff'), list):
                print('从市场获取股票数据时出错...')
                continue

            stock_data = res_dict['data']['diff']

            for item in stock_data:
                current_code = item['f12']

                if current_code in user_stock_codes:  # 检查是否已在正样本中
                    continue

                features.append([    # 不在正样本里，就添加进数据集，同时sustained必为0
                    item['f12'],
                    item['f2'],  # 注意这里是1297，网页上是12.97
                    item['f3'],  # 注意这里是2001，网页上是20.01%
                    item['f4'],  # 注意这里是282，网页上是2.82
                    item['f5'],  # 注意这里是637574，网页上是‘63.76万’
                    item['f7'],  # 注意这里是1732，网页上是17.32%
                    item['f15'],  # 注意这里是1691，网页上是16.91
                    item['f16'],  # 注意这里是1447，网页上是14.47
                    item['f17'],  # 注意这里是3000，网页上是30.00
                    item['f18'],  # 注意这里是1208，网页上是12.08
                    item['f9'],  # 注意这里是-7366，网页上是-73.66
                    item['f8']  # 注意这里是1784，网页上是17.84%
                ])
                target.append(0)  # 非持有标记
                negative_count += 1

        #print(f"已添加 {negative_count} 个负样本")

        return features, target

    def _preprocess_data(self, features, target):
        """
        数据预处理: 编码分类特征 + 数值特征转 float（处理非法字符如'-'）

        参数:
        features: 原始特征列表（每个元素是一个长度为12的列表）
        target: 目标列表

        返回:
        - X: 特征矩阵（所有数值特征 + 股票代码编码）
        - y: 目标向量
        """

        # 转为 DataFrame 便于处理
        columns = ['code', 'price', 'change1', 'change2', 'deal', 'range',
                   'high', 'low', 'open', 'close', 'PE', 'PB']
        df = pd.DataFrame(features, columns=columns)

        # Step 1: 股票代码编码（保留在第一列）
        df['code'] = self.label_encoder_code.fit_transform(df['code'])

        # Step 2: 将剩余的数值字段转换为 float（非法值转为 NaN，再统一填充为 0）
        numeric_cols = columns[1:]  # 除去 code
        df[numeric_cols] = df[numeric_cols].apply(
            pd.to_numeric, errors='coerce'
        ).fillna(0.0)

        # Step 3: 提取 X 和 y
        X = df.to_numpy()  # 转换为 NumPy 数组
        y = np.array(target)

        return X, y

    def train(self, user_id):
        """
        训练决策树模型

        返回: 训练后的模型
        """
        # 1. 加载数据
        features, target = self._load_data(user_id)

        # 2. 预处理数据
        self.X, self.y = self._preprocess_data(features, target)

        # 3. 训练模型
        self.tree.fit(self.X, self.y)

        return self.tree

    def predict_market_data(self, market_data):
        """
        使用已训练好的模型对市场数据进行预测。

        参数:
            market_data: List[Dict]，市场返回的数据，包含股票和基金

        返回:
            Tuple: (保留的基金列表, 预测为1的股票列表)
        """
        fund_list = []
        candidate_stocks = []
        stock_features = []
        stock_codes = []

        for item in market_data:
            code = item.get('股票代码') or item.get('基金代码')
            name = item.get('股票名称', '') or item.get('基金名称', '')
            type_hint = item.get('type', '')

            # 判断是否为基金
            if (type_hint == '基金') or name.endswith('基金') or (code and code.startswith(('5', '1'))):
                fund_list.append(item)
                continue

            # 是股票，则提取其特征用于预测
            try:
                stock_features.append([
                    item.get('股票代码','0'),
                    item.get('最新价', '0'),
                    item.get('涨跌幅', '0'),
                    item.get('涨跌额', '0'),
                    item.get('成交量', '0'),
                    item.get('振幅', '0'),
                    item.get('最高', '0'),
                    item.get('最低', '0'),
                    item.get('今开', '0'),
                    item.get('昨收', '0'),
                    item.get('市盈率', '0'),
                    item.get('换手率', '0')
                ])
                stock_codes.append(code)
                candidate_stocks.append(item)
            except Exception as e:
                print(f"提取特征出错，跳过该股票: {code}, 错误: {e}")
                continue

        if not stock_features:
            print("无可预测的股票")
            return fund_list, []


        # 用 DataFrame 来批量处理非法值，例如 '-' → NaN → 填 0
        df = pd.DataFrame(stock_features, columns=[
            'code', 'price', 'change1', 'change2', 'deal', 'range',
            'high', 'low', 'open', 'close', 'PE', 'PB'
        ])
        df = df.apply(pd.to_numeric, errors='coerce').fillna(0.0)
        df['code'] = self.label_encoder_code.fit_transform(df['code'])
        X_pred = df.to_numpy()

        y_pred = self.tree.predict(X_pred)

        # 只保留预测为 1 的股票
        selected_stocks = [
            stock for stock, pred in zip(candidate_stocks, y_pred) if pred == 1
        ]

        # 打印基金
        print("保留下来的基金：")
        for fund in fund_list:
            print(
                f"基金代码: {fund.get('基金代码', fund.get('股票代码'))}, 名称: {fund.get('基金名称', fund.get('股票名称'))}")

        # 打印预测为 1 的股票
        print("预测为 1 的股票：")
        for stock in selected_stocks:
            print(f"股票代码: {stock.get('股票代码')}, 名称: {stock.get('股票名称')}")

        return fund_list, selected_stocks



def main():
    #app = create_app()

    # app.config['SQLALCHEMY_DATABASE_URI'] = os.getenv('SQLALCHEMY_DATABASE_URI')
    # app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

    #with app.app_context():

        user_id = 1  # 实际使用时可以从JWT获取，这里先假设我们从外部获取到了 user_id

        # 创建训练器
        trainer = DecisionTreeTrainer()

        # 训练模型
        trainer.train(user_id)



if __name__ == "__main__":
    main()